# src/main.py

from math import log
from fastapi import FastAPI, HTTPException, BackgroundTasks, UploadFile, File
from pydantic import BaseModel
from typing import Optional
import uuid
import time
import httpx
import logging
from PIL import Image
import io
import os
import json
import platform
import psutil
from pathlib import Path
from datetime import datetime, timezone, timedelta
from cachetools import TTLCache # 导入缓存模块 (Import cache module)

# 从新模块导入AI服务 (Import AI service from the new module)
from .services.ai_services import get_ai_service, BaseAIService

# 导入日志配置 (Import logging configuration)
from .logging_config import setup_logging, log_violation

# --- Basic Setup (Keep as is) ---
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

with open('config.json', 'r') as f:
    config = json.load(f)

image_cache_dir = Path(config['image_service']['cache_dir'])
image_cache_dir.mkdir(exist_ok=True)

def print_startup_banner():
    """打印服务启动横幅信息"""
    try:
        # ASCII艺术字
        print("\n" + "="*80)
        print("""
██╗   ██╗██████╗ ██████╗  ██████╗    ██╗██╗  ██╗
██║   ██║██╔══██╗██╔══██╗██╔════╝   ██╔╝██║  ██║
██║   ██║██████╔╝██████╔╝██║       ██╔╝ ███████║
██║   ██║██╔══██╗██╔══██╗██║      ██╔╝  ╚════██║
╚██████╔╝██████╔╝██████╔╝╚██████╗██╔╝        ██║
 ╚═════╝ ╚═════╝ ╚═════╝  ╚═════╝╚═╝         ╚═╝
                                                                                                    
Universal Business Behavior Checker v4
        """)
        print("="*80)
        
        # 系统信息
        print("📊 系统信息")
        print("-" * 50)
        memory = psutil.virtual_memory()
        print(f"系统平台: {platform.system()} {platform.machine()} ({psutil.cpu_count()} cores)")
        print(f"内存使用: {memory.used // (1024*1024)}MB / {memory.total // (1024*1024)}MB")
        print(f"进程号  : {os.getpid()}")
        print(f"Python 版本  : Python {platform.python_version()}")
        print(f"启动时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print()
        
        # 服务信息
        print("🚀 服务信息")
        print("-" * 50)
        host = config['server']['host']
        port = config['server']['port']
        web_port = config['server']['web-port']
        ai_service = config['app_settings']['active_ai_service']
        print(f"API服务 : http://{host}:{port}")
        print(f"API端点 : http://{host}:{port}/api/v1/check-image")
        print(f"Web管理 : http://{host}:{web_port}")
        print(f"AI服务  : {ai_service}")
        print()
        
        # 操作提示
        print("💡 操作提示")
        print("-" * 50)
        print("• 按 Ctrl+C 优雅关闭服务")
        print(f"• 使用 POST http://{host}:{port}/api/v1/check-image 进行图片检查")
        print(f"• 访问 http://{host}:{web_port} 管理规则配置")
        print()
        
        print("✅ 服务已启动完成 - Powered by iSoftStone, 2025")
        print("="*80)
        print()
        
    except Exception as e:
        logger.warning(f"Failed to display startup banner: {e}")

app = FastAPI(title="图片内容检查服务", version="1.1") # Version bump

# --- Initialize Caching and AI Service ---
app_settings = config.get('app_settings', {})
analysis_cache = TTLCache(maxsize=app_settings.get('cache_max_size', 1000), ttl=app_settings.get('cache_ttl_seconds', 3600))
try:
    ai_service: BaseAIService = get_ai_service(config)
except ValueError as e:
    logger.fatal(f"AI Service initialization failed: {e}")
    # Exit if the core service cannot be loaded.
    exit(1)

# --- Model Definitions (Keep as is) ---
class ImageCheckRequest(BaseModel):
    image_url: str

class CheckResult(BaseModel):
    合规: bool
    说明: str
    详细检查结果: dict
    original_image_url: Optional[str] = None
    request_id: str
    timestamp: str

class UploadCheckResult(BaseModel):
    合规: bool
    说明: str
    详细检查结果: dict
    original_filename: Optional[str] = None
    request_id: str
    timestamp: str

# --- Rule Loading (Keep as is) ---
def load_rules():
    # ... (no changes to this function)
    try:
        with open('rule-definitions.json', 'r') as f:
            data = json.load(f)
            rules_dict = {}
            for rule in data['rules']:
                rules_dict[rule['rule_name']] = {
                    'keywords': rule['keywords'],
                    'description': rule['description']
                }
            return rules_dict
    except Exception as e:
        logger.error(f"加载规则失败: {str(e)}")
        return {
            "友商品牌": {"keywords": ["汽车之家", "易车"], "description": "检测图片中是否包含友商品牌名称"},
            "敏感内容": {"keywords": ["暴力", "色情", "赌博"], "description": "检测图片中是否包含敏感内容"}
        }
RULES = load_rules()

# --- Helper Functions (Refactored and New) ---
async def download_image(image_url: str) -> bytes:
    # ... (no changes to this function)
    try:
        async with httpx.AsyncClient(timeout=config['image_service']['timeout']) as client:
            response = await client.get(image_url)
            response.raise_for_status()
            return response.content
    except Exception as e:
        logger.error(f"下载图片失败: {str(e)}")
        raise HTTPException(status_code=400, detail=f"下载图片失败: {str(e)}")

async def process_image(image_bytes: bytes) -> bytes:
    """
    统一处理图片：转换格式并调整大小 (Process image: convert format and resize).
    """
    try:
        image = Image.open(io.BytesIO(image_bytes))
        
        # 转换WebP等格式为JPG (Convert WEBP etc. to JPG)
        if image.format == 'WEBP' or image.mode == 'RGBA':
            if image.mode == 'RGBA':
                background = Image.new('RGB', image.size, (255, 255, 255))
                background.paste(image, mask=image.split()[3])
                image = background
            else:
                image = image.convert('RGB')

        # 调整图片大小以优化性能 (Resize image for performance)
        max_size = config['image_service'].get('resize_max_size', 1024)
        if image.width > max_size or image.height > max_size:
            image.thumbnail((max_size, max_size))

        buffer = io.BytesIO()
        image.save(buffer, format="JPEG")
        return buffer.getvalue()
    except Exception as e:
        logger.error(f"图片处理失败 (format conversion/resize): {str(e)}")
        raise HTTPException(status_code=400, detail=f"图片处理失败: {str(e)}")

# REMOVED: analyze_image, convert_webp_to_jpg. Their logic is now in ai_services.py and process_image.

def extract_text_from_quotes(text: str) -> list:
    """
    从引号中提取文字内容
    支持单引号、双引号和各种引号格式
    """
    import re
    extracted_texts = []

    # 匹配各种引号格式
    quote_patterns = [
        r'"([^"]*)"',      # 英文双引号
        r"'([^']*)'",      # 英文单引号
        r'"([^"]*)"',      # 中文双引号
        r'"([^"]*)"',      # 中文单引号
        r'`([^`]*)`',      # 反引号
    ]

    for pattern in quote_patterns:
        matches = re.findall(pattern, text)
        extracted_texts.extend(matches)

    return extracted_texts

async def check_rules(image_analysis: str, ocr_text: str = "") -> dict:
    """
    检查规则违规情况
    增强版：同时检查AI分析结果、OCR识别文字和AI描述中的引号内容
    """
    details = []
    violated_rules = []
    detection_sources = {}  # 记录检测源信息

    # 1. 原始合并文本
    combined_text = f"{image_analysis} {ocr_text}".strip()

    # 2. 从AI分析结果中提取引号中的文字
    quoted_texts = extract_text_from_quotes(image_analysis)
    enhanced_text = combined_text

    # 将引号中的文字也加入到检查文本中
    if quoted_texts:
        quoted_combined = " ".join(quoted_texts)
        enhanced_text = f"{combined_text} {quoted_combined}".strip()

    logger.info(f"[规则检查] 增强文本长度: {len(enhanced_text)}")
    if quoted_texts:
        logger.info(f"[规则检查] 提取的引号文字: {quoted_texts}")

    for rule_name, rule_info in RULES.items():
        matched = False
        matched_keywords = []
        detection_info = {
            "ocr_text": False,
            "ai_description": False,
            "quoted_content": False
        }

        for keyword in rule_info['keywords']:
            # 检查原始OCR文字
            if keyword in ocr_text:
                if keyword not in matched_keywords:
                    matched = True
                    matched_keywords.append(keyword)
                    detection_info["ocr_text"] = True
                    logger.info(f"[规则检查] 在OCR文字中匹配关键词: '{keyword}'")

            # 检查AI分析结果
            if keyword in image_analysis:
                if keyword not in matched_keywords:
                    matched = True
                    matched_keywords.append(keyword)
                    detection_info["ai_description"] = True
                    logger.info(f"[规则检查] 在AI描述中匹配关键词: '{keyword}'")

            # 检查引号中的文字（包含子串匹配）
            if quoted_texts and any(keyword in quoted or any(keyword in word for word in quoted.split()) for quoted in quoted_texts):
                if keyword not in matched_keywords:
                    matched = True
                    matched_keywords.append(keyword)
                    detection_info["quoted_content"] = True
                    logger.info(f"[规则检查] 在引号内容中匹配关键词: '{keyword}'")

        if matched:
            violated_rules.append(rule_name)
            detection_sources[rule_name] = detection_info
            logger.info(f"[规则检查] 规则 '{rule_name}' 触发违规，匹配关键词: {matched_keywords}")

        details.append({
            "matched_conditions": matched_keywords,
            "rule_name": rule_name,
            "violated": matched,
            "detection_sources": detection_info
        })

    logger.info(f"[规则检查] 违规规则数量: {len(violated_rules)}")
    if violated_rules:
        logger.info(f"[规则检查] 违规规则: {', '.join(violated_rules)}")

    return {
        "details": details,
        "violated_rules": violated_rules,
        "detection_sources": detection_sources,
        "ocr_text": ocr_text,
        "quoted_texts": quoted_texts
    }

# --- API Endpoint (Refactored) ---
@app.post("/api/v1/check-image", response_model=CheckResult)
async def check_image(request: ImageCheckRequest, background_tasks: BackgroundTasks):
    """
    检查图片中是否包含违规内容 (Check image for non-compliant content).
    """
    request_id = str(uuid.uuid4())
    tz = timezone(timedelta(hours=8))
    timestamp = datetime.now(tz).strftime('%Y-%m-%dT%H:%M:%S+08:00')
    step_times = {}
    overall_start = time.time()

    # 步骤1: 检查缓存 (Step 1: Check cache)
    cached_result = analysis_cache.get(request.image_url)
    if cached_result:
        logger.info(f"Request ID: {request_id}, Cache HIT for URL: {request.image_url}")
        return cached_result

    logger.info(f"Request ID: {request_id}, Cache MISS for URL: {request.image_url}")

    try:
        # 步骤2: 下载和处理图片 (Step 2: Download and Process Image)
        start_time = time.time()
        image_bytes = await download_image(request.image_url)
        step_times['下载图片'] = time.time() - start_time

        start_time = time.time()
        processed_image_bytes = await process_image(image_bytes)
        step_times['图片预处理'] = time.time() - start_time

        # 步骤3: 获取OCR结果用于规则检查
        ocr_text = ""
        if hasattr(ai_service, 'ocr_service'):
            try:
                # 单独调用OCR服务获取原始识别结果
                ocr_result = await ai_service.ocr_service.analyze(processed_image_bytes)
                ocr_text = ocr_result.strip()
                logger.info(f"[规则检查] OCR识别文字: {ocr_text}")
            except Exception as e:
                logger.warning(f"[规则检查] OCR识别失败: {e}")

        # 步骤4: 调用AI服务分析 (Step 4: Call AI Service for Analysis)
        start_time = time.time()
        image_analysis = await ai_service.analyze(processed_image_bytes)
        ai_analysis_time = time.time() - start_time
        step_times['AI服务分析'] = ai_analysis_time

        # 记录AI服务决策信息 (Record AI service decision info)
        if hasattr(ai_service, 'ocr_service') and hasattr(ai_service, 'ai_service'):
            # 新混合服务模式 - 记录详细的决策过程
            service_decision = {
                'service_type': 'hybrid_v2',
                'active_model': getattr(ai_service, 'active_model', 'unknown'),
                'ai_analysis_time': ai_analysis_time,
                'total_processing_time': step_times.get('总耗时', 0),
                'strategy': 'conditional_ai_analysis',  # 有条件地使用AI分析
                'ocr_text_length': len(ocr_text)
            }
        else:
            # 单服务模式
            service_decision = {
                'service_type': config.get('active_model', config['app_settings'].get('active_ai_service', 'unknown')),
                'ai_analysis_time': ai_analysis_time,
                'strategy': 'direct_analysis'
            }

        step_times['service_decision'] = service_decision

        # 步骤5: 规则检查 (Step 5: Check Rules)
        start_time = time.time()
        check_result_details = await check_rules(image_analysis, ocr_text)
        step_times['规则检查'] = time.time() - start_time

        step_times['总耗时'] = time.time() - overall_start
        logger.info(f"Request ID: {request_id}, Total processing time: {step_times['总耗时']:.2f}s")
        check_result_details['step_times'] = step_times

        # 准备响应 (Prepare response)
        is_compliant = len(check_result_details['violated_rules']) == 0

        if is_compliant:
            description = "未发现违规内容"
        else:
            violated_rules = check_result_details['violated_rules']
            detection_sources = check_result_details.get('detection_sources', {})
            ocr_text = check_result_details.get('ocr_text', '')

            description_parts = []
            description_parts.append(f"在图片中发现违规内容: {', '.join(violated_rules)}")

            # # 添加检测源信息,这里的信息被我注释掉了，因为在生产会被显示出来.
            #if ocr_text:
            #   description_parts.append(f" OCR识别文字: {ocr_text}")
            if ocr_text:
               description_parts.append(f" {ocr_text}")      

            # 添加详细的检测源信息
            for rule_name, sources in detection_sources.items():
                source_info = []
                if sources.get('ocr_text'):
                    source_info.append("OCR文字")
                if sources.get('ai_description'):
                    source_info.append("AI描述")
                if sources.get('quoted_content'):
                    source_info.append("引号内容")

                if source_info:
                    description_parts.append(f"{rule_name}检测自: {', '.join(source_info)}")

            description = " | ".join(description_parts)

        response_data = {
            "合规": is_compliant,
            "说明": description,
            "详细检查结果": check_result_details,
            "original_image_url": request.image_url,
            "request_id": request_id,
            "timestamp": timestamp
        }

        logger.info(f"*** [RETURN to DongChenPeng] Response data: {response_data}")

        # 如果检测到违规内容，记录详细日志
        if not is_compliant:
            model_name = config.get('active_model', 'hybrid')
            log_violation(
                image_url=request.image_url,
                violation_reason=', '.join(check_result_details['violated_rules']),
                model_name=model_name,
                model_output=image_analysis
            )

        # 步骤5: 存储到缓存 (Step 5: Store in cache)
        analysis_cache[request.image_url] = response_data
        
        return response_data

    except HTTPException as e:
        raise e
    except Exception as e:
        logger.error(f"Request ID: {request_id}, unhandled exception: {str(e)}")
        raise HTTPException(status_code=500, detail=f"处理请求失败: {str(e)}")

@app.post("/api/v1/check-image-upload", response_model=UploadCheckResult)
async def check_image_upload(file: UploadFile = File(...), background_tasks: BackgroundTasks = None):
    """
    通过文件上传检查图片中是否包含违规内容 (Check uploaded image for non-compliant content).
    """
    request_id = str(uuid.uuid4())
    tz = timezone(timedelta(hours=8))
    timestamp = datetime.now(tz).strftime('%Y-%m-%dT%H:%M:%S+08:00')
    step_times = {}
    overall_start = time.time()

    try:
        # 步骤1: 验证文件类型
        if not file.content_type or not file.content_type.startswith('image/'):
            raise HTTPException(status_code=400, detail="请上传有效的图片文件")
        
        # 步骤2: 读取文件内容 (Step 2: Read file content)
        step_start = time.time()
        file_content = await file.read()
        if len(file_content) > config['image_service']['max_file_size']:
            raise HTTPException(status_code=400, detail="文件大小超过限制")
        step_times['文件读取'] = time.time() - step_start

        # 步骤3: 预处理图片 (Step 3: Preprocess image)
        step_start = time.time()
        processed_image_bytes = await process_image(file_content)
        step_times['图片预处理'] = time.time() - step_start

        # 步骤4: 获取OCR结果用于规则检查
        ocr_text = ""
        ai_service = get_ai_service(config)
        if hasattr(ai_service, 'ocr_service'):
            try:
                # 单独调用OCR服务获取原始识别结果
                ocr_result = await ai_service.ocr_service.analyze(processed_image_bytes)
                ocr_text = ocr_result.strip()
                logger.info(f"[规则检查] OCR识别文字: {ocr_text}")
            except Exception as e:
                logger.warning(f"[规则检查] OCR识别失败: {e}")

        # 步骤5: AI服务分析 (Step 5: AI service analysis)
        step_start = time.time()
        analyzed_text = await ai_service.analyze(processed_image_bytes)
        step_times['AI服务分析'] = time.time() - step_start

        # 步骤6: 规则检查 (Step 6: Rule checking)
        step_start = time.time()
        check_result_details = await check_rules(analyzed_text, ocr_text)
        step_times['规则检查'] = time.time() - step_start

        # 步骤7: 处理结果 (Step 7: Process results)
        step_times['总耗时'] = time.time() - overall_start

        # 记录AI服务决策信息
        ai_analysis_time = step_times['AI服务分析']
        if hasattr(ai_service, 'ocr_service') and hasattr(ai_service, 'ai_service'):
            # 新混合服务模式 - 记录详细的决策过程
            service_decision = {
                'service_type': 'hybrid_v2',
                'active_model': getattr(ai_service, 'active_model', 'unknown'),
                'ai_analysis_time': ai_analysis_time,
                'total_processing_time': step_times['总耗时'],
                'strategy': 'conditional_ai_analysis',  # 有条件地使用AI分析
                'ocr_text_length': len(ocr_text)
            }
        else:
            # 单服务模式
            service_decision = {
                'service_type': config.get('active_model', config['app_settings'].get('active_ai_service', 'unknown')),
                'ai_analysis_time': ai_analysis_time,
                'strategy': 'direct_analysis'
            }

        step_times['service_decision'] = service_decision
        check_result_details['step_times'] = step_times

        violated_rules = check_result_details.get('violated_rules', [])
        is_compliant = len(violated_rules) == 0

        if is_compliant:
            description = "未发现违规内容"
        else:
            detection_sources = check_result_details.get('detection_sources', {})
            ocr_text = check_result_details.get('ocr_text', '')

            description_parts = []
            description_parts.append(f"在图片中发现违规内容: {', '.join(violated_rules)}")

            # # 添加检测源信息,这里的信息被我注释掉了，因为在生产会被显示出来.
            #if ocr_text:
            #   description_parts.append(f" OCR识别文字: {ocr_text}")
            if ocr_text:
               description_parts.append(f" {ocr_text}")               

            # 添加详细的检测源信息
            for rule_name, sources in detection_sources.items():
                source_info = []
                if sources.get('ocr_text'):
                    source_info.append("OCR文字")
                if sources.get('ai_description'):
                    source_info.append("AI描述")
                if sources.get('quoted_content'):
                    source_info.append("引号内容")

                #这里的信息被我注释掉了，因为在生产会被显示出来.
                #if source_info:
                #    description_parts.append(f"{rule_name}检测自: {', '.join(source_info)}")

            description = " | ".join(description_parts)

        response_data = {
            "合规": is_compliant,
            "说明": description,
            "详细检查结果": check_result_details,
            "original_filename": file.filename,
            "request_id": request_id,
            "timestamp": timestamp
        }

        logger.info(f"*** [RETURN to DongChenPeng] Response data: {response_data}")


        # 如果检测到违规内容，记录详细日志
        if not is_compliant:
            model_name = config.get('active_model', 'hybrid')
            log_violation(
                image_url=f"上传文件: {file.filename}",
                violation_reason=', '.join(violated_rules),
                model_name=model_name,
                model_output=analyzed_text
            )
        
        return response_data

    except HTTPException as e:
        raise e
    except Exception as e:
        logger.error(f"Request ID: {request_id}, upload check exception: {str(e)}")
        raise HTTPException(status_code=500, detail=f"处理上传文件失败: {str(e)}")

# --- 启动事件处理 (Startup Event Handler) ---
@app.on_event("startup")
async def startup_event():
    """应用启动时的事件处理"""
    # 初始化日志系统
    loggers = setup_logging()
    logger.info("UBBC日志系统已初始化")

    print_startup_banner()

    # 启动Web管理界面
    import threading
    import uvicorn
    from .webui import app as web_app

    def start_web_admin():
        web_port = int(config['server']['web-port'])
        host = config['server']['host']
        print(f"🌐 启动Web管理界面: http://{host}:{web_port}")
        uvicorn.run(web_app, host=host, port=web_port)

    # 在后台线程启动Web管理界面
    web_thread = threading.Thread(target=start_web_admin, daemon=True)
    web_thread.start()

# --- Main Execution Block (Keep as is) ---
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host=config['server']['host'], port=int(config['server']['port']))
